import shutil
from keras.callbacks import Callback

from log import log
from util.model_listutil import ModelListUtil
from entity.train_data import TrainData
from predict import Predict
from plot import Plot
from dao.traindatadao import TrainDataDao
from dao.evaluationdao import EvaluationDao


class MyCallback(Callback):

    def on_train_end(self, logs=None):
        epoch_acc_img_name = f"storage/E-ACC{self.model.train_id}.png"
        Plot().plot_acc(self.model.history, epoch_acc_img_name)
        td = TrainDataDao()
        td.setImgNameCountEpochById(self.model.train_id, epoch_acc_img_name, self.model.current_epoch)

        max_acc_epoch = EvaluationDao().getMaxAccEpochByTrainId(self.model.train_id)
        log.critical(max_acc_epoch)
        shutil.copy('./checkpoints/' + str(max_acc_epoch) + 'checkpoints.h5', f'storage/{self.model.train_id}.h5')
        log.info("train_id"+str(self.model.train_id))
    def __init__(self, validation_data=()):
        super().__init__()
        self.x_test, self.y_test = validation_data
        self.train_data = TrainData()

    def on_train_begin(self, logs={}):
        self.model.current_epoch = 0
        t = TrainData()
        t.decay = self.model.conf.DECAY
        t.learning_rate = self.model.conf.LEARNING_RATE
        t.image_width = self.model.conf.IMAGE_WIDTH
        t.image_height = self.model.conf.IMAGE_HEIGHT
        t.count_epoch = self.model.current_epoch
        td = TrainDataDao()
        idx = td.setExcludeImgName(t, self.model.save_path)
        self.model.train_id = idx

        ModelListUtil().model_list_to_db(self.model.model_list, self.model.train_id)

    def on_epoch_end(self, epoch, logs={}):
        p = Predict()
        p.calc_metrics(self.model, self.x_test, self.y_test, epoch, self.model.train_id)
        self.model.save_weights('./checkpoints/' + str(self.model.current_epoch) + 'checkpoints.h5')

    def on_epoch_begin(self, epoch, logs=None):
        self.model.current_epoch += 1
